#include "SyntaxVisitors.h"
#include "CodeGenerator.h"
#include "BinaryCode.h"

namespace SimpleC
{
	namespace Compiler
	{
		void CodeGenerator::VisitProgram(ProgramSyntaxNode * program)
		{
			programSyntaxNode = program;
			ProgramBinaryCode = new ProgramCode();
			ProgramBinaryCode->StringBuffer.Add(L"");
			program->Functions.ForEach([&](RefPtr<FunctionSyntaxNode> f)
			{
				RefPtr<FunctionCode>functionBinaryCode = new FunctionCode();
				functionBinaryCode->Name = f->Name;
				ProgramBinaryCode->Functions.Add(functionBinaryCode);
				currentFunctionCode = functionBinaryCode.Ptr();
				f->Accept(this);
			});
		}

		void CodeGenerator::EmitInstruction(OpCode opCode)
		{
			Instruction instr;
			instr.IntOperand = 0;
			instr.Operator = opCode;
			currentFunctionCode->Instructions.Add(instr);
		}

		void CodeGenerator::EmitInstruction(OpCode opCode, int param)
		{
			Instruction instr;
			instr.IntOperand = param;
			instr.Operator = opCode;
			currentFunctionCode->Instructions.Add(instr);
		}

		void CodeGenerator::EmitInstruction(OpCode opCode, int param1, int param2)
		{
			Instruction instr;
			instr.Operand[0] = param1;
			instr.Operand[1] = param2;
			instr.Operator = opCode;
			currentFunctionCode->Instructions.Add(instr);
		}

		void CodeGenerator::EmitInstruction(OpCode opCode, double param)
		{
			Instruction instr;
			instr.DoubleOperand = param;
			instr.Operator = opCode;
			currentFunctionCode->Instructions.Add(instr);
		}

		int GetTypeSizeInBytes(ExpressionType & type)
		{
			switch (type.BaseType)
			{
			case ExpressionType::_Int:
				if (type.ArrayLength == 0)
					return 4;
				else
					return type.ArrayLength * 4;
				break;
			case ExpressionType::_Double:
				if (type.ArrayLength == 0)
					return 8;
				else
					return type.ArrayLength * 8;
				break;
			case ExpressionType::_Char:
				if (type.ArrayLength == 0)
					return 2;
				else
					return type.ArrayLength * 2;
				break;
			case ExpressionType::_String:
				if (type.ArrayLength == 0)
					return 4;
				else
					return type.ArrayLength * 4;
				break;
			default:
				return 0;
				break;
			}
		}

		void CodeGenerator::VisitFunction(FunctionSyntaxNode* function)
		{
			loopLabels.Clear();
			labelPosMap.Clear();
			currentVariables.Clear();
			int offset = -8;
			int localVarOffset = 4;
			paramSize = 0;
			for (int i = 0; i < function->Parameters.Count(); i++)
			{
				auto & var = function->Parameters[i];
				VariableOffset varOffset;
				varOffset.Type = var->Type->ToExpressionType();
				offset -= GetTypeSizeInBytes(varOffset.Type);
				paramSize += GetTypeSizeInBytes(varOffset.Type);
				varOffset.Offset = offset;
				currentVariables.Add(var->Name, _Move(varOffset));

				if (var->Type->ToExpressionType() == ExpressionType::String)
					currentFunctionCode->StrVarOffsets.Add(offset);
			}
			currentFunctionCode->ArgumentSize = paramSize;

			for (int i = 0; i < function->Variables.Count(); i++)
			{
				auto & var = function->Variables[i];
				VariableOffset varOffset;
				varOffset.Type = var.Type;
				
				varOffset.Offset = localVarOffset;
				localVarOffset += GetTypeSizeInBytes(varOffset.Type);
				currentVariables.Add(var.Name, _Move(varOffset));

				if (var.Type == ExpressionType::String)
					currentFunctionCode->StrVarOffsets.Add(varOffset.Offset);
			}
			localVarSize = localVarOffset - 4;
			

			if (function->IsExtern)
				currentFunctionCode->External = true;
			else
			{
				EmitInstruction(OpCode::PushEmpty, localVarSize);
				function->Body->Accept(this);
				EmitInstruction(OpCode::Ret, paramSize, 0);
				for (int i = 0; i < currentFunctionCode->Instructions.Count(); i++)
				{
					auto & instr = currentFunctionCode->Instructions[i];
					if (instr.Operator == OpCode::Jump ||
						instr.Operator == OpCode::JumpTrue ||
						instr.Operator == OpCode::JumpFalse)
					{
						instr.IntOperand = labelPosMap[instr.IntOperand];
					}
				}
			}
		}

		void CodeGenerator::VisitBlockStatement(BlockStatementSyntaxNode* stmt)
		{
			stmt->Statements.ForEach([&](RefPtr<StatementSyntaxNode> & node)
			{
				node->Accept(this);
			});
		}
		void CodeGenerator::VisitBreakStatement(BreakStatementSyntaxNode* stmt)
		{
			EmitInstruction(OpCode::Jump, GetCurrentLoopBreakLabel());
		}
		void CodeGenerator::VisitContinueStatement(ContinueStatementSyntaxNode* stmt)
		{
			EmitInstruction(OpCode::Jump,  GetCurrentLoopContinueLabel());
		}
		void CodeGenerator::VisitDoWhileStatement(DoWhileStatementSyntaxNode* stmt)
		{
			PushLoopLabels();
			SetCurrentLoopBeginPosition(GetCurrentInstructionPos());
			stmt->Statement->Accept(this);
			SetCurrentLoopContinuePosition(GetCurrentInstructionPos());
			stmt->Predicate->Accept(this);
			EmitInstruction(OpCode::JumpTrue, GetCurrentLoopBeginLabel());
			SetCurrentLoopBreakPosition(GetCurrentInstructionPos());
			PopLoopLabels();
		}
		void CodeGenerator::VisitEmptyStatement(EmptyStatementSyntaxNode* stmt){}
		void CodeGenerator::VisitForStatement(ForStatementSyntaxNode* stmt)
		{
			if (stmt->InitialExpression)
				stmt->InitialExpression->Accept(this);
			else if (stmt->VarDeclr)
			{
				stmt->VarDeclr->Accept(this);
			}
			PushLoopLabels();
			SetCurrentLoopBeginPosition(GetCurrentInstructionPos());

			if (stmt->MarginExpression)
			{
				stmt->MarginExpression->Accept(this);
				EmitInstruction(OpCode::JumpFalse, GetCurrentLoopBreakLabel());
			}
			stmt->Statement->Accept(this);
			SetCurrentLoopContinuePosition(GetCurrentInstructionPos());
			if (stmt->SideEffectExpression)
			{
				stmt->SideEffectExpression->Accept(this);
				EmitInstruction(OpCode::Pop, GetTypeSizeInBytes(stmt->SideEffectExpression->Type));
			}
			EmitInstruction(OpCode::Jump, GetCurrentLoopBeginLabel());
			SetCurrentLoopBreakPosition(GetCurrentInstructionPos());

			PopLoopLabels();
		}
		void CodeGenerator::VisitIfStatement(IfStatementSyntaxNode* stmt)
		{
			stmt->Predicate->Accept(this);
			int endLabel = labelPosMap.Count();
			labelPosMap.Add(-1);
			int elseLabel = labelPosMap.Count();
			labelPosMap.Add(-1);
			EmitInstruction(OpCode::JumpFalse, elseLabel);
			stmt->PositiveStatement->Accept(this);
			if (stmt->NegativeStatement)
			{
				EmitInstruction(OpCode::Jump, endLabel);
				labelPosMap[elseLabel] = GetCurrentInstructionPos();
				stmt->NegativeStatement->Accept(this);
			}
			else
				labelPosMap[elseLabel] = GetCurrentInstructionPos();
			labelPosMap[endLabel] = GetCurrentInstructionPos();
		}
		void CodeGenerator::VisitReturnStatement(ReturnStatementSyntaxNode* stmt)
		{
			int retSize = 0;
			if (stmt->Expression)
			{
				stmt->Expression->Accept(this);
				//register save return value..
				if (stmt->Expression->Type.BaseType != ExpressionType::_Void &&
					stmt->Expression->Type.BaseType != ExpressionType::_Function)
					retSize = GetTypeSizeInBytes(stmt->Expression->Type);
			}
			
			EmitInstruction(OpCode::Ret, paramSize, retSize);
		}
		void CodeGenerator::VisitVarDeclrStatement(VarDeclrStatementSyntaxNode* stmt)
		{
			for (int i = 0; i < stmt->Variables.Count(); i++)
			{
				auto & variable = stmt->Variables[i];
				if (variable->Expression)
				{
					variable->Expression->Accept(this);
					VariableOffset & varOffset = currentVariables[variable->Name];
					if (varOffset.Type == ExpressionType::Int)
						EmitInstruction(OpCode::iStore, varOffset.Offset);
					else if (varOffset.Type == ExpressionType::Char)
						EmitInstruction(OpCode::cStore, varOffset.Offset);
					else if (varOffset.Type == ExpressionType::Double)
						EmitInstruction(OpCode::dStore, varOffset.Offset);
					else if (varOffset.Type == ExpressionType::String)
						EmitInstruction(OpCode::sStore, varOffset.Offset);
					break;
				}
				else if (stmt->Type->ToExpressionType() == ExpressionType::String)
				{
					EmitInstruction(OpCode::sPush, 0);
					VariableOffset & varOffset = currentVariables[variable->Name];
					EmitInstruction(OpCode::sStore, varOffset.Offset);
				}
			}
		}
		void CodeGenerator::VisitWhileStatement(WhileStatementSyntaxNode* stmt)
		{
			PushLoopLabels();
			SetCurrentLoopBeginPosition(GetCurrentInstructionPos());
			SetCurrentLoopContinuePosition(GetCurrentInstructionPos());
			stmt->Predicate->Accept(this);
			EmitInstruction(OpCode::JumpFalse, GetCurrentLoopBreakLabel());
			stmt->Statement->Accept(this);
			EmitInstruction(OpCode::Jump, GetCurrentLoopBeginLabel());
			SetCurrentLoopBreakPosition(GetCurrentInstructionPos());
			PopLoopLabels();
		}
		void CodeGenerator::VisitExpressionStatement(ExpressionStatementSyntaxNode* stmt)
		{
			if (stmt->Expression)
				stmt->Expression->Accept(this);

			if (stmt->Expression->Type == ExpressionType::String)
				EmitInstruction(OpCode::sPop);
			else
			{
				if (stmt->Expression->Type.BaseType != ExpressionType::_Void)
					EmitInstruction(OpCode::Pop, GetTypeSizeInBytes(stmt->Expression->Type));
			}
		}

		void CodeGenerator::GenerateAssignment(ExpressionSyntaxNode * lvalueExpr)
		{
			if (dynamic_cast<VarExpressionSyntaxNode *>(lvalueExpr) != nullptr &&
				(	lvalueExpr->Type == ExpressionType::Int ||
					lvalueExpr->Type == ExpressionType::String ||
					lvalueExpr->Type == ExpressionType::Double ||
					lvalueExpr->Type == ExpressionType::Char))
			{
				int varOffset = GetVarOffset(((VarExpressionSyntaxNode*)lvalueExpr)->Variable);
				if (lvalueExpr->Type == ExpressionType::Int)
					EmitInstruction(OpCode::iStore, varOffset);
				else if (lvalueExpr->Type == ExpressionType::Double)
					EmitInstruction(OpCode::dStore, varOffset);
				else if (lvalueExpr->Type == ExpressionType::Char)
					EmitInstruction(OpCode::cStore, varOffset);
				else if (lvalueExpr->Type == ExpressionType::String)
					EmitInstruction(OpCode::sStore, varOffset);
			}
			else
			{
				GenerateLValueOffset(lvalueExpr);
				if (lvalueExpr->Type.BaseType == ExpressionType::_String)
					EmitInstruction(OpCode::sStoreOffset, GetTypeSizeInBytes(lvalueExpr->Type));
				else
					EmitInstruction(OpCode::StoreOffset, GetTypeSizeInBytes(lvalueExpr->Type));
			}
		}
		void CodeGenerator::VisitBinaryExpression(BinaryExpressionSyntaxNode* expr)
		{
			auto & leftType = expr->LeftExpression->Type;
			auto & rightType = expr->RightExpression->Type;
			auto & type = leftType;
			if (expr->Operator == Operator::Assign)
			{
				expr->RightExpression->Accept(this);
				if (leftType == ExpressionType::Int && rightType == ExpressionType::Double)
					EmitInstruction(OpCode::d2i);
				else if (leftType == ExpressionType::Double && rightType == ExpressionType::Int)
					EmitInstruction(OpCode::i2d);
				else if (leftType == ExpressionType::String && rightType == ExpressionType::Char)
					EmitInstruction(OpCode::c2s);

				GenerateAssignment(expr->LeftExpression.Ptr());
				expr->LeftExpression->Accept(this);
				return;
			}

			expr->LeftExpression->Accept(this);
			if (leftType == ExpressionType::Int && rightType == ExpressionType::Double)
			{
				EmitInstruction(OpCode::i2d);
				type = rightType;
			}
			if (leftType == ExpressionType::Char)
			{
				EmitInstruction(OpCode::c2s);
				type == ExpressionType::String;
			}
			expr->RightExpression->Accept(this);
			if (leftType == ExpressionType::Double && rightType == ExpressionType::Int)
				EmitInstruction(OpCode::i2d);
			if (rightType == ExpressionType::Char)
				EmitInstruction(OpCode::c2s);

			switch (expr->Operator)
			{
			case Operator::Add:
				if (type == ExpressionType::Int)
					EmitInstruction(OpCode::iAdd);
				else if (type == ExpressionType::Double)
					EmitInstruction(OpCode::dAdd);
				else if (type == ExpressionType::String)
					EmitInstruction(OpCode::sAdd);
				break;
			case Operator::Sub:
				if (type == ExpressionType::Int)
					EmitInstruction(OpCode::iSub);
				else if (type == ExpressionType::Double)
					EmitInstruction(OpCode::dSub);
				break;
			case Operator::Mul:
				if (type == ExpressionType::Int)
					EmitInstruction(OpCode::iMul);
				else if (type == ExpressionType::Double)
					EmitInstruction(OpCode::dMul);
				break;
			case Operator::Div:
				if (type == ExpressionType::Int)
					EmitInstruction(OpCode::iDiv);
				else if (type == ExpressionType::Double)
					EmitInstruction(OpCode::dDiv);
				break;
			case Operator::Mod:
				EmitInstruction(OpCode::Mod);
				break;
			case Operator::And:
				EmitInstruction(OpCode::And);
				break;
			case Operator::Or:
				EmitInstruction(OpCode::Or);
				break;
			case Operator::BitAnd:
				EmitInstruction(OpCode::BitAnd);
				break;
			case Operator::BitOr:
				EmitInstruction(OpCode::BitOr);
				break;
			case Operator::BitXor:
				EmitInstruction(OpCode::Xor);
				break;
			case Operator::Rsh:
				EmitInstruction(OpCode::Rsh);
				break;
			case Operator::Lsh:
				EmitInstruction(OpCode::Lsh);
				break;
			case Operator::Eql:
				if (type == ExpressionType::Int)
					EmitInstruction(OpCode::iEql);
				else if (type == ExpressionType::Double)
					EmitInstruction(OpCode::dEql);
				break;
			case Operator::Neq:
				if (type == ExpressionType::Int)
					EmitInstruction(OpCode::iNeq);
				else if (type == ExpressionType::Double)
					EmitInstruction(OpCode::dNeq);
				break;
			case Operator::Geq:
				if (type == ExpressionType::Int)
					EmitInstruction(OpCode::iGeq);
				else if (type == ExpressionType::Double)
					EmitInstruction(OpCode::dGeq);
				break;
			case Operator::Leq:
				if (type == ExpressionType::Int)
					EmitInstruction(OpCode::iLeq);
				else if (type == ExpressionType::Double)
					EmitInstruction(OpCode::dLeq);
				break;
			case Operator::Less:
				if (type == ExpressionType::Int)
					EmitInstruction(OpCode::iLess);
				else if (type == ExpressionType::Double)
					EmitInstruction(OpCode::dLess);
				break;
			case Operator::Greater:
				if (type == ExpressionType::Int)
					EmitInstruction(OpCode::iGreater);
				else if (type == ExpressionType::Double)
					EmitInstruction(OpCode::dGreater);
				break;
			default:
				break;
			}
		}
		void CodeGenerator::VisitConstantExpression(ConstantExpressionSyntaxNode* expr)
		{
			switch (expr->ConstType)
			{
			case ConstantExpressionSyntaxNode::ConstantType::Char:
				EmitInstruction(OpCode::cPush, expr->CharValue);
				break;
			case ConstantExpressionSyntaxNode::ConstantType::Int:
				EmitInstruction(OpCode::iPush, expr->IntValue);
				break;
			case ConstantExpressionSyntaxNode::ConstantType::String:
				EmitInstruction(OpCode::sPush, ProgramBinaryCode->StringBuffer.Count());
				ProgramBinaryCode->StringBuffer.Add(expr->StringValue);
				break;
			case ConstantExpressionSyntaxNode::ConstantType::Double:
				EmitInstruction(OpCode::dPush, expr->DoubleValue);
				break;
			default:
				break;
			}
		}
		void CodeGenerator::VisitIndexExpression(IndexExpressionSyntaxNode* expr)
		{
			GenerateLValueOffset(expr);
			int unitSize = GetTypeSizeInBytes(expr->Type) / max(1, expr->Type.ArrayLength);
			if (expr->Type.BaseType == ExpressionType::_String)
				EmitInstruction(OpCode::sLoadOffset, unitSize);
			else
				EmitInstruction(OpCode::LoadOffset, unitSize);
		}
		void CodeGenerator::VisitInvokeExpression(InvokeExpressionSyntaxNode* expr)
		{
			for (int i = expr->Arguments.Count() - 1; i >= 0; i--)
			{
				expr->Arguments[i]->Accept(this);
				if (expr->FunctionExpr->Type.Func->Parameters[i]->Type->ToExpressionType() == ExpressionType::Double &&
					expr->Arguments[i]->Type == ExpressionType::Int)
					EmitInstruction(OpCode::i2d);
				if (expr->FunctionExpr->Type.Func->Parameters[i]->Type->ToExpressionType() == ExpressionType::String &&
					expr->Arguments[i]->Type == ExpressionType::Char)
					EmitInstruction(OpCode::c2s);
			}
			expr->FunctionExpr->Accept(this);
			EmitInstruction(OpCode::Call);
		}

		void CodeGenerator::GenerateLValueOffset(ExpressionSyntaxNode * expr)
		{
			if (dynamic_cast<IndexExpressionSyntaxNode*>(expr) != 0)
			{
				auto idxExpr = (IndexExpressionSyntaxNode*)(expr);
				GenerateLValueOffset(idxExpr->BaseExpression.Ptr());
				idxExpr->IndexExpression->Accept(this);
				int unitSize = GetTypeSizeInBytes(expr->Type)/max(1, expr->Type.ArrayLength);
				EmitInstruction(OpCode::iPush, unitSize);
				EmitInstruction(OpCode::iMul);
				EmitInstruction(OpCode::iAdd);
			}
			else if (dynamic_cast<VarExpressionSyntaxNode*>(expr) != 0)
			{
				auto varExpr = (VarExpressionSyntaxNode *)(expr);
				EmitInstruction(OpCode::iPush, GetVarOffset(varExpr->Variable));
			}
			else
				throw "GenerateLValue Bug";
		}

		void CodeGenerator::VisitUnaryExpression(UnaryExpressionSyntaxNode* expr)
		{
			expr->Expression->Accept(this);

			switch (expr->Operator)
			{
			case Operator::Neg:
				if (expr->Expression->Type == ExpressionType::Int)
					EmitInstruction(OpCode::iNeg);
				else if (expr->Expression->Type == ExpressionType::Double)
					EmitInstruction(OpCode::dNeg);
				break;
			case Operator::Not:
				EmitInstruction(OpCode::Not);
				break;
			case Operator::PostDec:
				expr->Expression->Accept(this);
				EmitInstruction(OpCode::Dec);
				GenerateAssignment(expr->Expression.Ptr());
				break;
			case Operator::PostInc:
				expr->Expression->Accept(this);
				EmitInstruction(OpCode::Inc);
				GenerateAssignment(expr->Expression.Ptr());
				break;
			case Operator::PreDec:
				EmitInstruction(OpCode::Dec);
				GenerateAssignment(expr->Expression.Ptr());
				expr->Expression->Accept(this);
				break;
			case Operator::PreInc:
				EmitInstruction(OpCode::Inc);
				GenerateAssignment(expr->Expression.Ptr());
				expr->Expression->Accept(this);
				break;
			default:
				break;
			}
		}
		void CodeGenerator::VisitVarExpression(VarExpressionSyntaxNode* expr)
		{
			if(expr->Type.BaseType == ExpressionType::_Function)
			{
				int index = programSyntaxNode->Functions.BinarySearch(expr->Variable, [](RefPtr<FunctionSyntaxNode> & f, const String & str)->int{if (f->Name<str) return -1; else if (f->Name > str) return 1; else return 0;});
				if (index < 0)
					throw "Function does not exist.";
				EmitInstruction(OpCode::iPush, index);
			}
			else
			{
				int varOffset = GetVarOffset(expr->Variable);
				if (expr->Type == ExpressionType::Int)
					EmitInstruction(OpCode::iLoad, varOffset);
				else if (expr->Type == ExpressionType::Double)
					EmitInstruction(OpCode::dLoad, varOffset);
				else if (expr->Type == ExpressionType::Char)
					EmitInstruction(OpCode::cLoad, varOffset);
				else if (expr->Type == ExpressionType::String)
					EmitInstruction(OpCode::sLoad, varOffset);
				else
				{
					GenerateLValueOffset(expr);
					if (expr->Type.BaseType == ExpressionType::_String)
						EmitInstruction(OpCode::sLoadOffset, GetTypeSizeInBytes(expr->Type));
					else
						EmitInstruction(OpCode::LoadOffset, GetTypeSizeInBytes(expr->Type));
				}
			}
		}
		void CodeGenerator::VisitParameter(ParameterSyntaxNode* para){}
		void CodeGenerator::VisitType(TypeSyntaxNode* type){}
		void CodeGenerator::VisitDeclrVariable(VarDeclrStatementSyntaxNode::Variable* variable){}

		SyntaxVisitor * CreateCodeGenerator()
		{
			return new CodeGenerator();
		}
	}
}
